module collisions

  use dimensions

  implicit none

  complex, allocatable, private   ::   s(:,:,:,:,:)
  complex, allocatable, private   ::  s0(:,:,:,:,:)
 
contains


!-------------------------------------------------------------------------------------
subroutine evolve_F_collisions(dfu, dt)
!-------------------------------------------------------------------------------------

  use parameters, only: CollisionRate

  complex, intent(inout) :: dfu(jv, jz, jy, jx, jdir)
  real,    intent(in)    :: dt
  logical, save          :: first_call = .true.

  if (CollisionRate .eq. 0) return;

  if (first_call) then
     allocate(  s(jv, jz, jy, jx, jdir) )
     allocate( s0(jv, jz, jy, jx, jdir) )
     first_call = .false.
  end if

  s = dfu;
  call rk_rhs_collisions(s)
  s0 = s

  s = dfu + s*dt/2
  call rk_rhs_collisions(s)
  s0 = s0 + 2*s

  s = dfu + s*dt/2
  call rk_rhs_collisions(s)
  s0 = s0 + 2*s

  s = dfu + s*dt
  call rk_rhs_collisions(s)
  s0 = s0 + s

  dfu = dfu + s0*dt/6

end subroutine evolve_F_collisions


!-------------------------------------------------------------------------------------
subroutine rk_rhs_collisions(f)
!-------------------------------------------------------------------------------------

  use parameters, only: iunit, vPhase, CollisionRate
  use fft_wrap

  complex, intent(inout) :: f(jv, jz, jy, jx, jdir)
  integer                :: ix, iy, iz, iDir


  do iDir=1,jDir
     do ix=1, jx
        do iy=1, jy

           !$OMP PARALLEL PRIVATE(iz)
           !$OMP DO
           do iz=1, jz
              f(:,iz, iy, ix, iDir) = f(:, iz, iy, ix, iDir) * (vvec + vPhase) 
           end do
           !$OMP END DO
           !$OMP END PARALLEL
           
        end do
     enddo
  enddo


  call dfftw_execute_dft(plan_v5forward, f, f)
  f = f/jv

  do iDir=1,jDir
     do ix=1, jx
        do iy=1, jy
  
           !$OMP PARALLEL PRIVATE(iz)
           !$OMP DO
           do iz=1, jz
              f(:,iz, iy, ix, iDir) = f(:, iz, iy, ix, iDir) * iunit * kVvec(:) * CollisionRate
           end do
           !$OMP END DO
           !$OMP END PARALLEL

        end do
     enddo
  enddo


  call dfftw_execute_dft(plan_v5backward, f, f)


end subroutine rk_rhs_collisions

!-------------------------------------------------------------------------------------

end module collisions
